package com.nbsaas.boot.sql;

import java.util.*;
import java.util.regex.Pattern;

/**
 * SQL注入检查工具类
 * 用于检测SQL语句和参数中的潜在注入攻击
 */
public class SqlInjectionChecker {

    // SQL关键字黑名单（在参数中出现时才危险）
    private static final Set<String> DANGEROUS_KEYWORDS_IN_PARAMS = new HashSet<>(Arrays.asList(
            "UNION", "SELECT", "INSERT", "UPDATE", "DELETE", "DROP", "CREATE", "ALTER",
            "EXEC", "EXECUTE", "DECLARE", "TRUNCATE", "SCRIPT", "JAVASCRIPT", "VBSCRIPT",
            "ONLOAD", "ONERROR", "SCRIPT", "IFRAME", "OBJECT", "EMBED", "FORM",
            "XP_", "SP_", "OPENROWSET", "OPENDATASOURCE"
    ));

    // 危险的SQL操作符和函数
    private static final Set<String> DANGEROUS_OPERATORS = new HashSet<>(Arrays.asList(
            "--", "/*", "*/", "@@", "CHAR(", "ASCII(", "SUBSTRING(", "CAST(",
            "CONVERT(", "WAITFOR", "DELAY", "BENCHMARK(", "SLEEP("
    ));

    // 常见的注入攻击模式
    private static final List<Pattern> INJECTION_PATTERNS = Arrays.asList(
            // 单引号注入
            Pattern.compile(".*'.*[;|\\-\\-|#|/\\*].*", Pattern.CASE_INSENSITIVE),
            // 双引号注入
            Pattern.compile(".*\".*[;|\\-\\-|#|/\\*].*", Pattern.CASE_INSENSITIVE),
            // 注释注入
            Pattern.compile(".*(\\-\\-|\\/\\*|#).*", Pattern.CASE_INSENSITIVE),
            // UNION注入
            Pattern.compile(".*\\bUNION\\b.*\\bSELECT\\b.*", Pattern.CASE_INSENSITIVE),
            // 布尔盲注
            Pattern.compile(".*\\b(AND|OR)\\b.*\\b(TRUE|FALSE|1=1|1=2|0=0)\\b.*", Pattern.CASE_INSENSITIVE),
            // 时间盲注
            Pattern.compile(".*\\b(WAITFOR|DELAY|BENCHMARK|SLEEP)\\b.*", Pattern.CASE_INSENSITIVE),
            // 堆叠查询
            Pattern.compile(".*;\\s*(SELECT|INSERT|UPDATE|DELETE|DROP|CREATE|ALTER)\\b.*", Pattern.CASE_INSENSITIVE),
            // XSS相关
            Pattern.compile(".*<\\s*(script|iframe|object|embed|form)\\b.*", Pattern.CASE_INSENSITIVE),
            // 函数调用注入
            Pattern.compile(".*\\b(EXEC|EXECUTE|XP_|SP_)\\b.*", Pattern.CASE_INSENSITIVE)
    );

    /**
     * 检查SQL语句是否存在注入风险
     */
    public static SqlInjectionResult checkSql(String sql) {
        if (sql == null || sql.trim().isEmpty()) {
            return new SqlInjectionResult(false, "SQL语句为空");
        }

        String upperSql = sql.toUpperCase().trim();
        List<String> risks = new ArrayList<>();

        // 对于完整的SQL语句，只检查明显的注入模式，不检查正常的SQL关键字

        // 检查危险操作符
        for (String operator : DANGEROUS_OPERATORS) {
            if (sql.contains(operator)) {
                risks.add("检测到危险操作符: " + operator);
            }
        }

        // 检查注入模式
        for (Pattern pattern : INJECTION_PATTERNS) {
            if (pattern.matcher(sql).matches()) {
                risks.add("匹配到注入攻击模式: " + pattern.pattern());
            }
        }

        // 检查引号平衡
        if (!isQuotesBalanced(sql)) {
            risks.add("SQL语句中引号不平衡，可能存在注入风险");
        }

        // 检查分号数量（可能的堆叠查询）
        long semicolonCount = sql.chars().filter(ch -> ch == ';').count();
        if (semicolonCount > 1) {
            risks.add("检测到多个分号，可能存在堆叠查询注入");
        }

        boolean hasRisk = !risks.isEmpty();
        String message = hasRisk ? String.join("; ", risks) : "SQL语句安全";

        return new SqlInjectionResult(hasRisk, message, risks);
    }

    /**
     * 检查参数值是否存在注入风险
     */
    public static SqlInjectionResult checkParameter(Object parameter) {
        if (parameter == null) {
            return new SqlInjectionResult(false, "参数为空");
        }

        String paramStr = parameter.toString();
        String upperParam = paramStr.toUpperCase().trim();
        List<String> risks = new ArrayList<>();

        // 检查参数中的危险关键字
        for (String keyword : DANGEROUS_KEYWORDS_IN_PARAMS) {
            if (upperParam.contains(keyword)) {
                risks.add("参数中检测到危险关键字: " + keyword);
            }
        }

        // 检查危险操作符
        for (String operator : DANGEROUS_OPERATORS) {
            if (paramStr.contains(operator)) {
                risks.add("参数中检测到危险操作符: " + operator);
            }
        }

        // 检查注入模式
        for (Pattern pattern : INJECTION_PATTERNS) {
            if (pattern.matcher(paramStr).matches()) {
                risks.add("参数匹配到注入攻击模式: " + pattern.pattern());
            }
        }

        boolean hasRisk = !risks.isEmpty();
        String message = hasRisk ? String.join("; ", risks) : "参数安全";

        return new SqlInjectionResult(hasRisk, message, risks);
    }

    /**
     * 检查SQL条件列表是否存在注入风险
     */
    public static SqlInjectionResult checkConditions(List<SqlCondition> conditions) {
        if (conditions == null || conditions.isEmpty()) {
            return new SqlInjectionResult(false, "条件列表为空");
        }

        List<String> allRisks = new ArrayList<>();
        boolean hasAnyRisk = false;

        for (int i = 0; i < conditions.size(); i++) {
            SqlCondition condition = conditions.get(i);
            
            // 检查字段名
            SqlInjectionResult fieldResult = checkParameter(condition.getField());
            if (fieldResult.hasRisk()) {
                hasAnyRisk = true;
                allRisks.add("条件[" + i + "]字段名: " + fieldResult.getMessage());
            }

            // 检查操作符
            SqlInjectionResult operatorResult = checkParameter(condition.getOperator());
            if (operatorResult.hasRisk()) {
                hasAnyRisk = true;
                allRisks.add("条件[" + i + "]操作符: " + operatorResult.getMessage());
            }

            // 检查值
            SqlInjectionResult valueResult = checkParameter(condition.getValue());
            if (valueResult.hasRisk()) {
                hasAnyRisk = true;
                allRisks.add("条件[" + i + "]值: " + valueResult.getMessage());
            }

            // 检查第二个值（BETWEEN操作）
            if (condition.getValue2() != null) {
                SqlInjectionResult value2Result = checkParameter(condition.getValue2());
                if (value2Result.hasRisk()) {
                    hasAnyRisk = true;
                    allRisks.add("条件[" + i + "]第二个值: " + value2Result.getMessage());
                }
            }
        }

        String message = hasAnyRisk ? "检测到注入风险" : "所有条件安全";
        return new SqlInjectionResult(hasAnyRisk, message, allRisks);
    }

    /**
     * 安全的SQL拼接方法，在拼接前进行注入检查
     */
    public static String safeAppend(String sql, List<SqlCondition> conditions) throws Exception {
        // 检查原始SQL
        SqlInjectionResult sqlResult = checkSql(sql);
        if (sqlResult.hasRisk()) {
            throw new SqlInjectionException("原始SQL存在注入风险: " + sqlResult.getMessage());
        }

        // 检查条件
        SqlInjectionResult conditionsResult = checkConditions(conditions);
        if (conditionsResult.hasRisk()) {
            throw new SqlInjectionException("SQL条件存在注入风险: " + conditionsResult.getMessage());
        }

        // 如果检查通过，则进行正常的SQL拼接
        String result = SqlWhereAppender.append(sql, conditions);
        
        // 对最终结果再次检查
        SqlInjectionResult finalResult = checkSql(result);
        if (finalResult.hasRisk()) {
            throw new SqlInjectionException("最终SQL存在注入风险: " + finalResult.getMessage());
        }

        return result;
    }

    /**
     * 检查引号是否平衡
     */
    private static boolean isQuotesBalanced(String sql) {
        int singleQuoteCount = 0;
        int doubleQuoteCount = 0;
        boolean escaped = false;

        for (int i = 0; i < sql.length(); i++) {
            char c = sql.charAt(i);
            
            if (escaped) {
                escaped = false;
                continue;
            }
            
            if (c == '\\') {
                escaped = true;
                continue;
            }
            
            if (c == '\'') {
                singleQuoteCount++;
            } else if (c == '"') {
                doubleQuoteCount++;
            }
        }

        return singleQuoteCount % 2 == 0 && doubleQuoteCount % 2 == 0;
    }

    /**
     * SQL注入检查结果类
     */
    public static class SqlInjectionResult {
        private final boolean hasRisk;
        private final String message;
        private final List<String> riskDetails;

        public SqlInjectionResult(boolean hasRisk, String message) {
            this(hasRisk, message, new ArrayList<>());
        }

        public SqlInjectionResult(boolean hasRisk, String message, List<String> riskDetails) {
            this.hasRisk = hasRisk;
            this.message = message;
            this.riskDetails = riskDetails != null ? riskDetails : new ArrayList<>();
        }

        public boolean hasRisk() {
            return hasRisk;
        }

        public String getMessage() {
            return message;
        }

        public List<String> getRiskDetails() {
            return riskDetails;
        }

        @Override
        public String toString() {
            return "SqlInjectionResult{" +
                    "hasRisk=" + hasRisk +
                    ", message='" + message + '\'' +
                    ", riskDetails=" + riskDetails +
                    '}';
        }
    }

    /**
     * SQL注入异常类
     */
    public static class SqlInjectionException extends RuntimeException {
        public SqlInjectionException(String message) {
            super(message);
        }

        public SqlInjectionException(String message, Throwable cause) {
            super(message, cause);
        }
    }
} 